from domain.model.platform.base import BasePlatformChatModel


class ZhipuChatModel(BasePlatformChatModel):
    PLATFORM_NAME = "zhipu"
    SUPPORT_MODEL_LIST = ['glm-4', 'glm-4-plus', 'glm-4-long', 'glm-4-flash-250414']
    MODEL_INSTANCE_MAP = {}

    def init_model(self, *model_names: str, **options: object) -> "BasePlatformChatModel":
        super().init_model(*model_names, **options)
        from langchain_community.chat_models import ChatZhipuAI
        for model in model_names:
            self.MODEL_INSTANCE_MAP[model] = ChatZhipuAI(
                model=model,
                zhipuai_api_key=options.get('api_key'),
                temperature=options.get('temperature', 0),
                max_tokens=options.get('max_token', 8192),
            )
        return self


zhipu_chat_model = ZhipuChatModel()
